{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "O9I9rz0pTamX"
   },
   "source": [
    "# XNLI Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EiowR0WNTd1C"
   },
   "source": [
    "In this notebook, we are going to train a model for evaluating on [XNLI](https://github.com/facebookresearch/XNLI). XNLI a cross-lingual NLI task, spanning 15 different languages, with 10,000 validation and test examples per language. Notably, XNLI does not have its own training set - instead, the usual recipe is to MNLI as a training set, and is then zero-shot evaluated on NLI examples in other languages. Of course, this works best when you start with a model that has already been pretrained on a lot of multi-lingual text, such as mBERT or XLM/XLM-RoBERTa.\n",
    "\n",
    "Hence, the tricky part about this setup is that although we have separate XNLI and MNLI tasks, we need them to all use the same task head. We will cover how to easily do this with `jiant`.\n",
    "\n",
    "--- \n",
    "\n",
    "In this notebook, we will:\n",
    "\n",
    "* Train an XLM-RoBERTa base model on MNLI\n",
    "* Evaluate on XNLI-de (German) and XNLI-zh (Chinese)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rXbD_U1_VDnw"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tC9teoazUnW8"
   },
   "source": [
    "#### Install dependencies\n",
    "\n",
    "First, we will install libraries we need for this code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8aU3Z9szuMU9"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!git clone https://github.com/nyu-mll/jiant.git\n",
    "%cd jiant\n",
    "!pip install -r requirements-no-torch.txt\n",
    "!pip install --no-deps -e ./"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KGJcCmRzU1Qb"
   },
   "source": [
    "#### Download data\n",
    "\n",
    "Next, we will download MNLI and XNLI data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jKCz8VksvFlN"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%cd /content\n",
    "# Download MNLI and XNLI data\n",
    "!PYTHONPATH=/content/jiant python jiant/jiant/scripts/download_data/runscript.py \\\n",
    "    download \\\n",
    "    --tasks mnli xnli \\\n",
    "    --output_path=/content/tasks/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rQKSAhYzVIlv"
   },
   "source": [
    "## `jiant` Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "v88oXqmBvFuK"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/content/jiant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ibmMT7CXv1_P"
   },
   "outputs": [],
   "source": [
    "import jiant.proj.main.tokenize_and_cache as tokenize_and_cache\n",
    "import jiant.proj.main.export_model as export_model\n",
    "import jiant.proj.main.scripts.configurator as configurator\n",
    "import jiant.proj.main.runscript as main_runscript\n",
    "import jiant.shared.caching as caching\n",
    "import jiant.utils.python.io as py_io\n",
    "import jiant.utils.display as display\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HPZHyLOlVp07"
   },
   "source": [
    "#### Download model\n",
    "\n",
    "Next, we will download an `xlm-roberta-base` model. This also includes the tokenizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "K06qUGjkKWa7"
   },
   "outputs": [],
   "source": [
    "export_model.export_model(\n",
    "    hf_pretrained_model_name_or_path=\"xlm-roberta-base\",\n",
    "    output_base_path=\"./models/xlm-roberta-base\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dV-T-8r1V0wf"
   },
   "source": [
    "#### Tokenize and cache\n",
    "\n",
    "With the model and data ready, we can now tokenize and cache the inputs features for our tasks. This converts the input examples to tokenized features ready to be consumed by the model, and saved them to disk in chunks.\n",
    "\n",
    "Note that we are tokenize `train` and `val` data for MNLI, but only `val` data for XNLI, since there is no corresponding training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "22bNWQajO4zm"
   },
   "outputs": [],
   "source": [
    "# Tokenize and cache MNLI\n",
    "tokenize_and_cache.main(tokenize_and_cache.RunConfiguration(\n",
    "    task_config_path=f\"./tasks/configs/mnli_config.json\",\n",
    "    hf_pretrained_model_name_or_path=\"xlm-roberta-base\",\n",
    "    output_dir=f\"./cache/mnli\",\n",
    "    phases=[\"train\", \"val\"],\n",
    "))\n",
    "\n",
    "# Tokenize and cache XNLI-de, XNLI-zh\n",
    "for lang in [\"de\", \"zh\"]:\n",
    "    tokenize_and_cache.main(tokenize_and_cache.RunConfiguration(\n",
    "        task_config_path=f\"./tasks/configs/xnli_{lang}_config.json\",\n",
    "        hf_pretrained_model_name_or_path=\"xlm-roberta-base\",\n",
    "        output_dir=f\"./cache/xnli_{lang}\",\n",
    "        phases=[\"val\"],\n",
    "    ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JJ-mWSQQWJsw"
   },
   "source": [
    "We can inspect the first examples of the first chunk of each task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iLk_X0KypUyr"
   },
   "outputs": [],
   "source": [
    "row = caching.ChunkedFilesDataCache(\"./cache/mnli/train\").load_chunk(0)[0][\"data_row\"]\n",
    "print(row.input_ids)\n",
    "print(row.tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2n00e6Xrp1bI"
   },
   "outputs": [],
   "source": [
    "row = caching.ChunkedFilesDataCache(\"./cache/xnli_de/val\").load_chunk(0)[0][\"data_row\"]\n",
    "print(row.input_ids)\n",
    "print(row.tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "A7FxZgEbqCx-"
   },
   "outputs": [],
   "source": [
    "row = caching.ChunkedFilesDataCache(\"./cache/xnli_zh/val\").load_chunk(0)[0][\"data_row\"]\n",
    "print(row.input_ids)\n",
    "print(row.tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3MBuH19IWOr0"
   },
   "source": [
    "#### Writing a run config\n",
    "\n",
    "Here we are going to write what we call a `jiant_task_container_config`. This configuration file basically defines a lot of the subtleties of our training pipeline, such as what tasks we will train on, do evaluation on, batch size for each task. The new version of `jiant` leans heavily toward explicitly specifying everything, for the purpose of inspectability and leaving minimal surprises for the user, even as the cost of being more verbose.\n",
    "\n",
    "We use a helper \"Configurator\" to write out a `jiant_task_container_config`, since most of our setup is pretty standard. We specify to train only on MNLI, but evaluate on MNLI, XNLI-de and XNLI-zh.\n",
    "\n",
    "**Depending on what GPU your Colab session is assigned to, you may need to lower the train batch size.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pQYtl7xTKsiP"
   },
   "outputs": [],
   "source": [
    "jiant_run_config = configurator.SimpleAPIMultiTaskConfigurator(\n",
    "    task_config_base_path=\"./tasks/configs\",\n",
    "    task_cache_base_path=\"./cache\",\n",
    "    train_task_name_list=[\"mnli\"],\n",
    "    val_task_name_list=[\"mnli\", \"xnli_de\", \"xnli_zh\"],\n",
    "    train_batch_size=32,\n",
    "    eval_batch_size=64,\n",
    "    epochs=0.1,\n",
    "    num_gpus=1,\n",
    ").create_config()\n",
    "display.show_json(jiant_run_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-UF501yoXHBi"
   },
   "source": [
    "To briefly go over the major components of the `jiant_task_container_config`:\n",
    "\n",
    "* `task_config_path_dict`: The paths to the task config files we wrote above.\n",
    "* `task_cache_config_dict`: The paths to the task features caches we generated above.\n",
    "* `sampler_config`: Determines how to sample from different tasks during training.\n",
    "* `global_train_config`: The number of total steps and warmup steps during training.\n",
    "* `task_specific_configs_dict`: Task-specific arguments for each task, such as training batch size and gradient accumulation steps.\n",
    "* `taskmodels_config`: Task-model specific arguments for each task-model, including what tasks use which model.\n",
    "* `metric_aggregator_config`: Determines how to weight/aggregate the metrics across multiple tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sf9C_RlBX-lK"
   },
   "source": [
    "**We need to make one small change to the auto-generated config**: we need to ensure that all three tasks use the same model head. Otherwise, each task will have its own task head, and the XNLI heads will be untrained.\n",
    "\n",
    "We can make a simple change to the dictionary, setting all of them to point to an `nli_model` head, and then write out the config."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "InlkHBiCTXT9"
   },
   "outputs": [],
   "source": [
    "jiant_run_config[\"taskmodels_config\"][\"task_to_taskmodel_map\"] = {\n",
    "    \"mnli\": \"nli_model\",\n",
    "    \"xnli_de\": \"nli_model\",\n",
    "    \"xnli_zh\": \"nli_model\",\n",
    "}\n",
    "os.makedirs(\"./run_configs/\", exist_ok=True)\n",
    "py_io.write_json(jiant_run_config, \"./run_configs/jiant_run_config.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BBKkvXzdYPqZ"
   },
   "source": [
    "#### Start training\n",
    "\n",
    "Finally, we can start our training run. \n",
    "\n",
    "Before starting training, the script also prints out the list of parameters in our model. You should notice that the only task head is the `nli_model` head."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JdwWPgjQWx6I"
   },
   "outputs": [],
   "source": [
    "run_args = main_runscript.RunConfiguration(\n",
    "    jiant_task_container_config_path=\"./run_configs/jiant_run_config.json\",\n",
    "    output_dir=\"./runs/run1\",\n",
    "    hf_pretrained_model_name_or_path=\"xlm-roberta-base\",\n",
    "    model_path=\"./models/xlm-roberta-base/model/model.p\",\n",
    "    model_config_path=\"./models/xlm-roberta-base/model/config.json\",\n",
    "    learning_rate=1e-5,\n",
    "    eval_every_steps=500,\n",
    "    do_train=True,\n",
    "    do_val=True,\n",
    "    force_overwrite=True,\n",
    ")\n",
    "\n",
    "main_runscript.run_loop(run_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4SXcuHFIYp6Y"
   },
   "source": [
    "Finally, we should see the validation scores for MNLI, XNLI-de, and XNLI-zh. Given that the training data is in English, we expect to see slightly higher scores for MNLI, but the scores for XNLI-de and XNLI-zh are still decent!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "jiant XNLI Example",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}